modelopt/torch/export/plugins/hf_spec_export.py (4)

16-16: Fix typos in module docstring.

"Modifiy stated_dict" → "Modify state_dict".

-"""Modifiy stated_dict and config for exporting speculative decoding in official format."""
+"""Modify state_dict and config for exporting speculative decoding in official format."""

---

23-24: Remove or use unused SPECULATIVE_DECODING_MODES.

Currently unused; either hook it into the guards or drop it to avoid dead code.

-SPECULATIVE_DECODING_MODES = ["eagle", "medusa"]
+# Reserved for future multi-mode handling:
+# SPECULATIVE_DECODING_MODES = ["eagle", "medusa"]

---

25-44: Rename constant to fix typo and improve clarity.

EALGE_MODELOPT_TO_OFFICIAL → EAGLE_MODELOPT_TO_OFFICIAL; update references.

-EALGE_MODELOPT_TO_OFFICIAL = {
+EAGLE_MODELOPT_TO_OFFICIAL = {
@@
-    _check_state_dict_keys_match(model.eagle_module, EALGE_MODELOPT_TO_OFFICIAL["required"])
+    _check_state_dict_keys_match(model.eagle_module, EAGLE_MODELOPT_TO_OFFICIAL["required"])
@@
-        **EALGE_MODELOPT_TO_OFFICIAL["required"],
-        **EALGE_MODELOPT_TO_OFFICIAL["optional"],
+        **EAGLE_MODELOPT_TO_OFFICIAL["required"],
+        **EAGLE_MODELOPT_TO_OFFICIAL["optional"],

Also applies to: 65-66, 70-72

---

68-79: Avoid repeated state_dict() calls and use provided post_state_dict for fallback.

Slight perf/clarity win and safer fallback when lm_head lives only in the passed-in state dict.

-    export_state_dict = {}
-    for ours_key, export_key in {
-        **EALGE_MODELOPT_TO_OFFICIAL["required"],
-        **EALGE_MODELOPT_TO_OFFICIAL["optional"],
-    }.items():
-        if ours_key in model.eagle_module.state_dict():
-            export_state_dict[export_key] = model.eagle_module.state_dict()[ours_key]
+    module_state = model.eagle_module.state_dict()
+    export_state_dict = {}
+    for ours_key, export_key in {
+        **EAGLE_MODELOPT_TO_OFFICIAL["required"],
+        **EAGLE_MODELOPT_TO_OFFICIAL["optional"],
+    }.items():
+        if ours_key in module_state:
+            export_state_dict[export_key] = module_state[ours_key]
@@
-    if "eagle_lm_head.weight" not in model.eagle_module.state_dict():
-        export_state_dict["lm_head.weight"] = model.state_dict()["lm_head.weight"]
+    if "eagle_lm_head.weight" not in module_state:
+        if "lm_head.weight" in post_state_dict:
+            export_state_dict["lm_head.weight"] = post_state_dict["lm_head.weight"]
+        else:
+            # Fall back to model.state_dict() if needed
+            base_state = model.state_dict()
+            if "lm_head.weight" in base_state:
+                export_state_dict["lm_head.weight"] = base_state["lm_head.weight"]
+            else:
+                raise KeyError("lm_head.weight not found in post_state_dict or model.state_dict()")

modelopt/torch/speculative/plugins/transformers.py (4)

721-727: Align reverse_mapping device with logits to avoid cross-device indexing.

Safer if full_logits is on a different device than d2t.

-        reverse_mapping = (
-            torch.arange(len(self.eagle_module.d2t)).to(self.eagle_module.d2t.device)
-            + self.eagle_module.d2t
-        )
+        reverse_mapping = (
+            torch.arange(len(self.eagle_module.d2t), device=full_logits.device, dtype=torch.long)
+            + self.eagle_module.d2t.to(full_logits.device)
+        )
         return full_logits[:, :, reverse_mapping]

---

856-858: Nit: fix comment typo.

"diabled" → "disabled".

-            # NOTE: diabled for now.
+            # NOTE: disabled for now.

---

1052-1058: Compute masked accuracy over valid positions only.

Current mean divides by B*T; use sum over mask / mask.sum() for accurate reporting when masks zero-out prefixes.

-        base_predict_tok = base_model_logits.argmax(dim=-1)
-        eagle_predict_tok = eagle_logits.argmax(dim=-1)
-        accuracy = (
-            (loss_mask[:, :, 0] * (base_predict_tok == eagle_predict_tok)).float().mean().item()
-        )
-        accuracy = round(accuracy, 3)
+        base_predict_tok = base_model_logits.argmax(dim=-1)
+        eagle_predict_tok = eagle_logits.argmax(dim=-1)
+        valid = loss_mask[:, :, 0].bool()
+        correct = (base_predict_tok == eagle_predict_tok) & valid
+        denom = valid.sum().clamp_min(1).float()
+        accuracy = round(correct.sum().float().div(denom).item(), 3)

---

716-717: Add explicit d2t buffer assertion before vocab remap
Although gating on draft_vocab_size != vocab_size is correct, insert an assertion (e.g. assert hasattr(self, "d2t"), "d2t buffer not initialized") immediately before calling _map_logits_to_draft_vocab to surface misconfigurations more clearly.

modelopt/torch/export/plugins/__init__.py (1)

18-24: Optional: make public surface explicit.

Consider defining all in hf_spec_export to avoid star-import drift.

 with import_plugin("transformers"):
-    from .hf_spec_export import *
+    from .hf_spec_export import *  # relies on hf_spec_export.__all__

examples/speculative_decoding/export_hf_checkpoint.py (2)

25-29: Make CLI args required and self-documenting

Avoid silent defaults. Require both paths and add help for better UX.

-def parse_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model_path", type=str, default="")
-    parser.add_argument("--export_path", type=str, default="")
-    return parser.parse_args()
+def parse_args():
+    parser = argparse.ArgumentParser(description="Export a HF checkpoint (with ModelOpt state) for deployment.")
+    parser.add_argument("--model_path", type=str, required=True, help="Path or HF hub id of the trained checkpoint.")
+    parser.add_argument("--export_path", type=str, required=True, help="Destination directory for exported files.")
+    return parser.parse_args()

---

34-41: Set eval mode prior to export

Ensures deterministic modules (e.g., dropout) are disabled during the dummy forward used in export.

 args = parse_args()
-model = AutoModelForCausalLM.from_pretrained(args.model_path, torch_dtype="auto")
+model = AutoModelForCausalLM.from_pretrained(args.model_path, torch_dtype="auto")
+model.eval()
 with torch.inference_mode():
     export_hf_checkpoint(
         model,  # The quantized model.
         export_dir=args.export_path,  # The directory where the exported files will be stored.
     )

modelopt/torch/export/unified_export_hf.py (2)

497-503: Make _quant_applied resilient to missing keys

Minor polish: use dict.get to avoid KeyError if the structure changes.

-def _quant_applied(hf_quant_config: dict) -> bool:
-    """Check if any quantization is applied."""
-    return not (
-        hf_quant_config["quantization"]["quant_algo"] == QUANTIZATION_NONE
-        and not hf_quant_config["quantization"]["quantized_layers"]
-    )
+def _quant_applied(hf_quant_config: dict) -> bool:
+    """Check if any quantization is applied."""
+    q = hf_quant_config.get("quantization", {})
+    return not (q.get("quant_algo") == QUANTIZATION_NONE and not q.get("quantized_layers"))

---

30-33: Optional: avoid hard import if plugins are unused

If plugins are only relevant for Eagle exports, consider importing inside export_hf_checkpoint to reduce import-time side effects and avoid potential cycles.

examples/speculative_decoding/train_eagle3_and_export.sh (4)

53-53: Fix ShellCheck SC2155: avoid command substitution in export

Assign first, then export to prevent masking return values.

-  export CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((NUM_GPU-1)))
+  devs="$(seq -s, 0 $((NUM_GPU-1)))"
+  export CUDA_VISIBLE_DEVICES="$devs"

---

58-66: Create output dirs proactively

Avoid relying on downstream scripts to mkdir. Harmless if they already exist.

 echo "==== [1/3] Training draft model ===="
 OUTPUT_DIR=ckpts/${MODEL_BASENAME}-$(date +%Y%m%d_%H%M)
+mkdir -p "$(dirname "$OUTPUT_DIR")"
 ./launch_train.sh --model $BASE_MODEL \
             --output_dir $OUTPUT_DIR \
             --data $DATA \
             --num_gpu $NUM_GPU \
             --num_epochs 2 \
             --eagle_config eagle_config.json

---

70-72: Also mkdir for export path

Prevents failures if parent dir is missing.

 echo "==== [3/3] Exporting checkpoint to deployment format ===="
 EXPORT_PATH=export/${MODEL_BASENAME}-$(date +%Y%m%d_%H%M)
+mkdir -p "$(dirname "$EXPORT_PATH")"
 python export_hf_checkpoint.py --model_path $OUTPUT_DIR --export_path $EXPORT_PATH

---

25-41: Minor: align flag name in comment with implementation

The comment says --base-model but code uses --base_model.

examples/speculative_decoding/README.md (11)

7-9: Fix grammar and proper names in intro.

Use articles and brand casing; current phrasing is awkward.

-This folder contains end-to-end runnable speculative decoding fine-tuning pipeline where Llama3.2-1B from huggingface is trained on Daring-Anteater dataset.
-
-This example focus on training with HF. To train with Megatron-LM, please refer to [this link](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/post_training/modelopt) in Megatron-LM repo.
+This folder contains an end-to-end runnable speculative decoding fine‑tuning pipeline in which Llama‑3.2‑1B (Hugging Face) is trained on the Daring‑Anteater dataset.
+
+This example focuses on training with Hugging Face. To train with Megatron‑LM, see the [Megatron‑LM example](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/post_training/modelopt).

---

65-71: Avoid suggesting hardcoded API keys in examples.

Remove the fake key or clearly denote an env var. This reduces the chance of users pasting secrets.

-pip install vllm
-vllm serve meta-llama/Llama-3.2-1B-Instruct --api-key token-abc123 --port 8000  --tensor-parallel-size 1
+pip install vllm
+# If your deployment requires an API key, set VLLM_API_KEY in the environment instead of hardcoding.
+VLLM_API_KEY=... vllm serve meta-llama/Llama-3.2-1B-Instruct --port 8000 --tensor-parallel-size 1

---

90-96: Tighten wording and fix minor grammar.

-For eagle1 and eagle3 we provide an [default model architecture config](.../default_config.py#L18) in modelopt. User can overwrite default settings by providing additional json dict. In this example, we overwrite the `draft_vocab_size` by in `eagle_config.json`:
+For EAGLE‑1 and EAGLE‑3 we provide a [default model architecture config](.../default_config.py#L18) in ModelOpt. You can override default settings by providing an additional JSON dict. In this example, we override `draft_vocab_size` in `eagle_config.json`:

---

100-108: Fix typos (“tokenzier”, “hugginface”) and add clarity.

-`main.py` provides a example for converting a base HF model for speculative decoding and training it. It consists of a few simple steps:
-First, load base model and tokenzier from hugginface:
+`main.py` provides an example for converting a base HF model for speculative decoding and training it. It consists of a few simple steps:
+First, load the base model and tokenizer from Hugging Face:

---

118-121: Typo in key name placeholder.

-# overwrite config with custom config
-config["eagle_architecture_config"].update({"<overwrite_kyes>": "<overwrite_values>"})
+# Overwrite config with custom config
+config["eagle_architecture_config"].update({"<overwrite_keys>": "<overwrite_values>"})

---

131-136: Fix typo and verify API alias.

Spelling: “deocoding” → “decoding”. Also, please confirm mtsp.convert is the correct symbol/alias in this repo.

-Then, we convert model to a speculative deocoding model:
+Then, we convert the model to a speculative decoding model:

---

140-147: Avoid private Trainer APIs; rely on Trainer for device placement.

trainer._move_model_to_device is a private method and may change; Trainer already handles device placement. Also confirm the correct path for enable_huggingface_checkpointing().

-trainer._move_model_to_device(model, trainer.args.device)
-
 # Enable HF checkpointing so that the saved model will contain the speculative decoding module
-mto.enable_huggingface_checkpointing()
+from modelopt.torch import export as mto_export  # adjust import as needed
+mto_export.enable_huggingface_checkpointing()

---

152-161: Fix spelling and tighten phrasing in launch instructions.

-... along with a bash script to launch the training with huggingface accelrate in `launch_train.sh`, which can be runned by:
+... along with a bash script to launch training with Hugging Face Accelerate in `launch_train.sh`, which can be run by:

---

189-197: Minor YAML lead‑in punctuation and naming consistency.

-To serve the checkpoint with trtllm, we can run trtllm-serve with:
+To serve the checkpoint with TRT‑LLM, run `trtllm-serve`:
 ...
-,
- with `extra-llm-api-config.yml` being
+with `extra-llm-api-config.yml`:

---

224-233: Brand/style fixes in Support Matrix.

-| LLAMA 2 | ✅ | ✅ | ✅ |
-| LLAMA 3, 3.1 | ✅ | ✅ | ✅ |
+| Llama 2 | ✅ | ✅ | ✅ |
+| Llama 3, 3.1 | ✅ | ✅ | ✅ |
 ...
-| QWen 1.5,2,2.5 | ✅ | ✅ | ✅ |
+| Qwen 1.5, 2, 2.5 | ✅ | ✅ | ✅ |

---

236-238: Use correct company casing (“NVIDIA”).

-Ready-to-deploy speculation module checkpoints \[[🤗 Hugging Face - Nvidia TensorRT Model Optimizer Collection](https://huggingface.co/collections/nvidia/model-optimizer-66aa84f7966b3150262481a4)\]
+Ready-to-deploy speculation module checkpoints \[[🤗 Hugging Face - NVIDIA TensorRT Model Optimizer Collection](https://huggingface.co/collections/nvidia/model-optimizer-66aa84f7966b3150262481a4)\]

examples/speculative_decoding/eagle_config.json (1)

2-8: Align rope_scaling schema with expected HF conventions (consider adding "type": "dynamic").

Given the presence of low_freq_factor, high_freq_factor, and original_max_position_embeddings, many HF configs expect rope_scaling.type="dynamic". If the exporter relies on HF-compatible fields, add the type or confirm your plugin translates this correctly.

Apply this minimal diff if compatible with your pipeline:

   "rope_scaling": {
+        "type": "dynamic",
         "factor": 32.0,
         "low_freq_factor": 1.0,
         "high_freq_factor": 4.0,
         "original_max_position_embeddings": 8192,
         "rope_type": "llama3"
   },

examples/speculative_decoding/server_generate.py (1)

56-56: Make --system_prompt truly optional to avoid odd joining behavior.

With nargs="+" and default "", " ".join(args.system_prompt) can behave unexpectedly. Prefer a list default.

Apply:

-parser.add_argument("--system_prompt", nargs="+", type=str, default="", help="System prompt")
+parser.add_argument("--system_prompt", nargs="*", type=str, default=[], help="System prompt")

system_prompt = " ".join(args.system_prompt) will remain correct (empty string when not provided).

Also applies to: 185-186

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

examples/speculative_decoding/launch_train.sh (4)

95-97: Sanitize OUTPUT_DIR model basename to avoid filesystem issues.

Some model names contain characters problematic for paths. Sanitize the basename before composing OUTPUT_DIR.

 MODEL_BASENAME=$(basename "$MODEL")
-OUTPUT_DIR=${OUTPUT_DIR:-"ckpts/${MODEL_BASENAME}-$(date +%Y%m%d_%H%M)"}
+MODEL_BASENAME_SAFE="${MODEL_BASENAME//[^a-zA-Z0-9._-]/_}"
+OUTPUT_DIR=${OUTPUT_DIR:-"ckpts/${MODEL_BASENAME_SAFE}-$(date +%Y%m%d_%H%M)"}

---

77-83: Don’t hard-fail on unknown flags; forward them to the training script.

Exiting on first unknown arg prevents users from passing through Hugging Face/Accelerate options. Collect and forward instead.

-    *)
-      >&2 printf "Error: Invalid argument ${1#*=}\n"
-      exit 1
-      ;;
+    *)
+      EXTRA_ARGS+=("$1")
+      ;;

Add initialization near the top (before the while-loop):

 set -eo pipefail
+
+EXTRA_ARGS=()

---

130-153: Plumb through extra args to CMD.

Forward any collected unknown args to main.py.

     --data_path $DATA \
-    $SPECULATIVE_ARGS
+    $SPECULATIVE_ARGS \
+    ${EXTRA_ARGS[@]:-}

---

141-142: Make --do_eval overridable while defaulting to False.

Hardcoding --do_eval False blocks standalone use; keep False by default but allow override via a parsed --do_eval.

-    --do_eval False \
+    --do_eval ${DO_EVAL:-False} \

Add a parser case (outside this hunk) to honor --do_eval:

# In the case/esac parser:
--do_eval*)
  if [[ "$1" != *=* ]]; then shift; fi
  DO_EVAL="${1#*=}"
  ;;

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

modelopt/torch/export/unified_export_hf.py (1)

529-545: Guard spec-decoding transforms when ModelOpt state is absent

Calling spec-decoding transforms unconditionally can break vanilla HF exports. Add a safe gate.

-        post_state_dict = rename_and_prune_if_spec_decoding(model, post_state_dict)
+        if hasattr(model, "_modelopt_state") and getattr(model, "_modelopt_state") is not None:
+            post_state_dict = rename_and_prune_if_spec_decoding(model, post_state_dict)
@@
-        config_data = set_config_if_spec_decoding(model, config_data)
+        if hasattr(model, "_modelopt_state") and getattr(model, "_modelopt_state") is not None:
+            config_data = set_config_if_spec_decoding(model, config_data)

examples/speculative_decoding/README.md (3)

19-20: Broken anchors from prior commit are now correct.
“Complete Workflow” and “Support Matrix” Jump To targets are fixed.

---

172-174: Code-fence language fixed.
CLI is now fenced as bash, not python.

---

181-185: Export section fences and copy now correct.
Language tag and “format” typo are fixed.

examples/speculative_decoding/SLURM_prepare_data.md (1)

9-10: Quote job name; normalize spacing

Avoid parsing issues with ':' in job name; trim double spaces.

-salloc  -N4 -A <account> -p <partition>  -J <account>-synthetic:data-gen -t 120
+salloc -N4 -A <account> -p <partition> -J "<account>-synthetic:data-gen" -t 120

modelopt/torch/export/unified_export_hf.py (1)

517-526: Path handling nit: prefer Path operations over f-strings

Improves readability and OS-compatibility.

-        with open(f"{export_dir}/hf_quant_config.json", "w") as file:
+        with open(Path(export_dir) / "hf_quant_config.json", "w") as file:
             json.dump(hf_quant_config, file, indent=4)
@@
-        original_config = f"{export_dir}/config.json"
+        original_config = Path(export_dir) / "config.json"
         config_data = {}
@@
-        with open(original_config, "w") as file:
+        with open(original_config, "w") as file:
             json.dump(config_data, file, indent=4)

Also applies to: 536-547

modelopt/torch/speculative/plugins/transformers.py (4)

689-717: Consistent draft-vocab remap policy

_base path remaps logits only during training; teacher-provided path remaps unconditionally. Ensure both paths align with how downstream loss expects logits to be in draft space.

-        if self.eagle_config.draft_vocab_size != self.eagle_config.vocab_size and self.training:
+        if self.eagle_config.draft_vocab_size != self.eagle_config.vocab_size:
             assert hasattr(self.eagle_module, "d2t"), "d2t buffer not initialized"
             base_model_logits = self._map_logits_to_draft_vocab(base_model_logits)

If inference truly needs full-vocab logits, gate at call site instead.

---

720-726: Bounds/correctness checks and caching for mapping

Add sanity checks and cache the reverse index to avoid recomputing per call.

-    def _map_logits_to_draft_vocab(self, full_logits):
-        reverse_mapping = (
-            torch.arange(len(self.eagle_module.d2t)).to(self.eagle_module.d2t.device)
-            + self.eagle_module.d2t
-        )
-        return full_logits[:, :, reverse_mapping]
+    def _map_logits_to_draft_vocab(self, full_logits):
+        d2t = self.eagle_module.d2t
+        draft = d2t.numel()
+        if not hasattr(self, "_reverse_draft_index") or self._reverse_draft_index.numel() != draft:
+            self._reverse_draft_index = torch.arange(draft, device=d2t.device) + d2t
+        # Assert indices are in range
+        assert self._reverse_draft_index.max().item() < full_logits.size(-1), "draft→full index OOB"
+        return full_logits.index_select(-1, self._reverse_draft_index)

---

826-831: Aux hidden states: validate presence/shape in teacher path

If use_aux_hidden_state is True and teacher path is used, assert aux_hidden_states is provided and has expected last-dim = len(layer_ids)*hidden_size to fail fast with a clear message.

-                if "base_model_outputs" in kwargs:
-                    aux_hidden_states = kwargs["base_model_outputs"]["aux_hidden_states"]
+                if "base_model_outputs" in kwargs:
+                    aux_hidden_states = kwargs["base_model_outputs"].get("aux_hidden_states")
+                    assert aux_hidden_states is not None, "aux_hidden_states required for EAGLE-3 teacher path"

---

790-803: Clarify DynamicCache initialization and document aux_hidden_states requirement
The past_key_values = None branch is already handled by DynamicCache.from_legacy_cache(None), so no extra guard is required; consider adding an inline comment (# will initialize a fresh DynamicCache below) for clarity. Update the EAGLE-3 docs to specify that when use_aux_hidden_state is enabled, base_model_outputs must include an aux_hidden_states tensor shaped to match eagle_aux_hidden_state_layer_ids.

examples/speculative_decoding/README.md (10)

28-28: Use consistent capitalization: “ModelOpt”.
Brand is capitalized elsewhere; align here.

-Install Modelopt with `hf` dependencies and other requirements for this example:
+Install ModelOpt with `hf` dependencies and other requirements for this example:

---

47-47: Capitalize “ModelOpt”.

-This one-line command runs a minimal example workflow of training and exporting an EAGLE draft model in Modelopt.
+This one-line command runs a minimal example workflow of training and exporting an EAGLE draft model in ModelOpt.

---

83-90: Explain d2t mapping with a concrete relation.
Add a short sentence showing how the mapping composes tokens.

-This will produce a `d2t.pt` file in `save_dir`, which is the mapping from draft vocabs to full vocab that will be read by our draft model later.
+This produces a `d2t.pt` file in `save_dir`, a mapping from draft vocab to full vocab used by the draft model. Conceptually: `tokens_full = tokens_draft + d2t[tokens_draft]`.

---

101-101: Capitalize “ModelOpt”.

-### Training Draft Model with Modelopt
+### Training Draft Model with ModelOpt

---

106-110: Import missing dependency in snippet.
Readers will copy/paste this; include the import.

-```python
-model = transformers.AutoModelForCausalLM.from_pretrained(
+```python
+import transformers
+model = transformers.AutoModelForCausalLM.from_pretrained(

---

166-170: Capitalize “ModelOpt”.

-The saved modelopt checkpoint is similar in architecture to HF models. It can be further optimized through **ModelOpt**, e.g., PTQ and QAT.
+The saved ModelOpt checkpoint is similar in architecture to HF models. It can be further optimized through **ModelOpt**, e.g., PTQ and QAT.

-After training draft model, we can evaluate the saved modelopt checkpoint on MT-bench by:
+After training the draft model, evaluate the saved ModelOpt checkpoint on MT‑Bench with:

---

227-235: Polish model names and add “last updated”.
Use canonical casing and note when the matrix was last validated.

-| LLAMA 2 | ✅ | ✅ | ✅ |
-| LLAMA 3, 3.1 | ✅ | ✅ | ✅ |
+| Llama 2 | ✅ | ✅ | ✅ |
+| Llama 3, 3.1 | ✅ | ✅ | ✅ |
 | Mistral | ✅ | ✅ | ✅ |
 | Phi 3 | ✅ | ✅ | ✅ |
-| QWen 1.5,2,2.5 | ✅ | ✅ | ✅ |
+| Qwen 1.5, 2, 2.5 | ✅ | ✅ | ✅ |

Optionally add a line above the table: “Last validated: 2025‑09‑05.”

---

93-99: Clarify EAGLE-1 vs EAGLE-3 default configs and update link

• In the README example, call out which settings enable EAGLE-3 (use_aux_hidden_state, eagle_aux_hidden_state_layer_ids, use_mtp_layernorm, eagle_disable_moe, eagle_hidden_state_distillation, etc.) and update the URL to point directly at the EagleConfig default JSON in the ModelOpt API docs.

---

65-71: vLLM flags and quantization option are up-to-date

• --api-key (or VLLM_API_KEY) is the current auth flag.
• --quantization=modelopt has been supported since v0.6.5.

Consider showing environment-based auth (VLLM_API_KEY/BASE_URL) in examples instead of embedding tokens.

---

194-215: Update README example to match current trtllm-serve flags and YAML schema

• Confirm the CLI flag --extra_llm_api_options is still supported in your installed version (it appears in the latest docs).
• Under speculative_config, use the up-to-date keys: speculative_decoding_mode (alias decoding_type), max_draft_len, speculative_model_dir, etc.
• Under kv_cache_config, include valid fields such as enable_block_reuse, enable_partial_reuse, copy_on_partial_reuse, free_gpu_memory_fraction, and any other keys required by your target release.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

modelopt/torch/export/unified_export_hf.py (2)

544-544: Also guard spec-config transform to avoid AttributeError.

Mirror the pruning guard for config transform.

(Implemented in the diff above.)

---

529-530: Guard spec-decoding state-dict pruning (works on vanilla HF models).

Call pruning only when ModelOpt state exists.

-        post_state_dict = rename_and_prune_if_spec_decoding(model, post_state_dict)
+        if hasattr(model, "_modelopt_state"):
+            post_state_dict = rename_and_prune_if_spec_decoding(model, post_state_dict)

modelopt/torch/export/plugins/hf_spec_export.py (2)

87-95: Guard _modelopt_state access to avoid crashes on vanilla HF.

Directly indexing model._modelopt_state can raise AttributeError. Use the same guard pattern as in rename_and_prune_if_spec_decoding.

-def set_config_if_spec_decoding(model: nn.Module, config_data: dict):
+def set_config_if_spec_decoding(model: nn.Module, config_data: dict):
@@
-    if len(model._modelopt_state) != 1 or model._modelopt_state[0][0] != "eagle":
-        # return as is
-        return config_data
+    opt_modes = getattr(model, "_modelopt_state", None)
+    if (
+        not isinstance(opt_modes, (list, tuple))
+        or len(opt_modes) != 1
+        or opt_modes[0][0] != "eagle"
+    ):
+        return config_data

---

95-155: Merge with existing config instead of overwriting it.

Dropping unknown keys can remove fields like quantization_config; merge template into incoming config and deep-merge eagle_config.

-    return template_config
+    # Merge: preserve existing fields while overriding with official template values.
+    merged = {**config_data, **template_config}
+    merged["eagle_config"] = {
+        **config_data.get("eagle_config", {}),
+        **template_config["eagle_config"],
+    }
+    return merged

examples/speculative_decoding/SLURM_prepare_data.md (4)

3-5: Tighten wording and fix grammar.

Minor clarity nits.

-For basic parallelization of synthetic data generation we provide some SLURM support.
-Assuming a `$SLURM_JOB_ID` is present and nodes, n1, n2, n3, n4 are selected the following is achievable.
+For basic parallelization of synthetic data generation, we provide SLURM support.
+Assuming `$SLURM_JOB_ID` is present and nodes n1, n2, n3, n4 are allocated, the following is achievable.

---

12-16: Polish sentence casing.

-Create shards of some given size
+Create shards of a given size

---

24-26: Capitalize framework names.

-This will launch a vllm server (sglang is also available) on each node. Each node will work through 10 shards of data (10*max_lines_per_shard number of samples).
+This launches a vLLM server (SGLang is also available) on each node. Each node will work through 10 shards of data (10*max_lines_per_shard samples).

---

29-31: Keep argument list consistent between runs.

The second launch example omits the system prompt argument; include it or add a note that it’s optional to avoid confusion.

-bash distributed_generate/launch.sh $SLURM_JOB_ID vllm TinyLlama/TinyLlama-1.1B-Chat-v1.0 /data/train/ /data/output /scripts/ 40 10 n1,n2,n3,n4
+bash distributed_generate/launch.sh $SLURM_JOB_ID vllm TinyLlama/TinyLlama-1.1B-Chat-v1.0 /data/train/ /data/output /scripts/ 40 10 n1,n2,n3,n4 "\"You are a helpful assistant.\""

examples/speculative_decoding/README.md (5)

18-18: Capitalize EAGLE and tighten phrasing.

-| Simplified Workflow | Train, evaluate, and export eagle model with one-line command | \[[Link](#getting-started-simplified-workflow)\] |
+| Simplified Workflow | Train, evaluate, and export the EAGLE model with a one‑line command | \[[Link](#getting-started-simplified-workflow)\] |

---

35-39: Add Git LFS note for dataset clone.

HF datasets via git require LFS; add a one‑liner to reduce user friction.

 ```bash
+git lfs install
 git clone https://huggingface.co/datasets/nvidia/Daring-Anteater

---

`69-69`: **Clarify note reads better as a flag tip.**



```diff
-Note: Add `--quantization=modelopt` flag for quantized models.
+Tip: Add `--quantization=modelopt` when serving quantized models.

---

101-101: Consistent branding: ModelOpt.

-### (Optional) Configuring Draft Model
+### (Optional) Configuring Draft Model (ModelOpt)

---

227-235: Model names casing and minor formatting.

-| LLAMA 2 | ✅ | ✅ | ✅ |
-| LLAMA 3, 3.1 | ✅ | ✅ | ✅ |
+| Llama 2 | ✅ | ✅ | ✅ |
+| Llama 3, 3.1 | ✅ | ✅ | ✅ |

modelopt/torch/speculative/plugins/transformers.py (1)

713-717: Guard d2t usage with a clearer error.

Assertion is fine, but a ValueError is more actionable when draft vocab remapping is expected.

-            assert hasattr(self.eagle_module, "d2t"), "d2t buffer not initialized"
+            if not hasattr(self.eagle_module, "d2t"):
+                raise ValueError("Draft‑vocab remapping requested but eagle_module.d2t is not initialized")

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.359Z
Learning: In the TensorRT-Model-Optimizer codebase, both EAGLE1_DEFAULT_CFG and EAGLE3_DEFAULT_CFG in config.py already use deepcopied configurations (either directly via deepcopy() or through variables that were already deepcopied), making additional deepcopy calls unnecessary when using ALGO_TO_CONFIG mapping in tests.

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.359Z
Learning: In the TensorRT-Model-Optimizer codebase, the EAGLE1_DEFAULT_CFG and EAGLE3_DEFAULT_CFG configurations in config.py are already deepcopied, so additional deepcopy calls are not needed when using ALGO_TO_CONFIG mapping in tests.

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.359Z
Learning: In the TensorRT-Model-Optimizer codebase, EAGLE1_DEFAULT_CFG in config.py already uses deepcopy(default_eagle_config) in its definition, making additional deepcopy calls unnecessary when using ALGO_TO_CONFIG mapping in tests.

examples/speculative_decoding/README.md (2)

15-22: TOC anchors now resolve correctly.

Broken anchors/typo from earlier review are fixed.

---

172-185: Code fences now correctly use bash and grammar is fixed.

Prior issues are resolved.

examples/speculative_decoding/README.md (4)

49-53: Avoid line-number deep links that can drift.

Linking to ...default_config.py#L18 is brittle. Point to the file (or a permalink) instead.

-- Initializes the draft model with [default settings](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/modelopt/torch/speculative/eagle/default_config.py#L18)
+- Initializes the draft model with [default settings](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/modelopt/torch/speculative/eagle/default_config.py)

---

83-90: Great fix on the flag; tighten wording and clarify mapping.

The extra dashes issue is resolved. Suggest minor wording + add a one-liner showing how d2t is applied.

-We can optionally use smaller vocab size for the draft model for faster training and inference. E.g. Llama3.2-1B has a vocab size of 128256. In this example, we construct a draft vocab mapping of size 32k by finding the most commonly appeared vocabs in our training set:
+We can optionally use a smaller draft vocabulary to speed up training/inference. For example, Llama‑3.2‑1B uses a 128,256‑token vocabulary. Here we build a 32k draft‑to‑target mapping by selecting the most frequent tokens in the training set:
@@
-This will produce a `d2t.pt` file in `save_dir`, which is the mapping from draft vocabs to full vocab that will be read by our draft model later.
+This produces `d2t.pt` in `save_dir`, a draft‑to‑target mapping the draft model will load later (applied roughly as: accepted_tokens = draft_tokens + d2t[draft_tokens]).

---

93-99: Link stability.

config.py#L37 may drift; consider linking to the file (or a permalink).

---

166-167: Unify brand capitalization: ModelOpt.

Use “ModelOpt” consistently.

-The saved modelopt checkpoint is similar in architecture to HF models. It can be further optimized through **ModelOpt**, e.g., PTQ and QAT.
+The saved ModelOpt checkpoint is similar in architecture to HF models. It can be further optimized through **ModelOpt**, e.g., PTQ and QAT.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.359Z
Learning: In the TensorRT-Model-Optimizer codebase, both EAGLE1_DEFAULT_CFG and EAGLE3_DEFAULT_CFG in config.py already use deepcopied configurations (either directly via deepcopy() or through variables that were already deepcopied), making additional deepcopy calls unnecessary when using ALGO_TO_CONFIG mapping in tests.

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.359Z
Learning: In the TensorRT-Model-Optimizer codebase, the EAGLE1_DEFAULT_CFG and EAGLE3_DEFAULT_CFG configurations in config.py are already deepcopied, so additional deepcopy calls are not needed when using ALGO_TO_CONFIG mapping in tests.

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.359Z
Learning: In the TensorRT-Model-Optimizer codebase, EAGLE1_DEFAULT_CFG in config.py already uses deepcopy(default_eagle_config) in its definition, making additional deepcopy calls unnecessary when using ALGO_TO_CONFIG mapping in tests.

examples/speculative_decoding/launch_train.sh (1)

87-91: Guard torch detection and prevent division-by-zero; prefer NUM_GPU when available.

python -c 'import torch' can fail under set -e, and torch.cuda.device_count() may return 0, yielding a division-by-zero or DEFAULT_SAVE_STEPS=0. Use ${NUM_GPU} when valid, fall back robustly to 1, and clamp to min 1.

-# Get the default value for save_steps based on the available number of GPUs
-GPU_COUNT=$(python -c "import torch; print(torch.cuda.device_count())")
-# Calculate save_steps
-DEFAULT_SAVE_STEPS=$((8192 / GPU_COUNT))
+# Get the default value for save_steps based on the available number of GPUs
+# Prefer NUM_GPU if provided; else try torch; finally fall back to 1.
+if [[ -n "${NUM_GPU:-}" && "${NUM_GPU}" =~ ^[0-9]+$ && ${NUM_GPU} -gt 0 ]]; then
+  GPU_COUNT=${NUM_GPU}
+else
+  GPU_COUNT=$(python - <<'PY'
+try:
+  import torch
+  print(torch.cuda.device_count() or 0)
+except Exception:
+  print(0)
+PY
+  ) || GPU_COUNT=0
+  [[ -z "$GPU_COUNT" || "$GPU_COUNT" -le 0 ]] && GPU_COUNT=1
+fi
+# Calculate save_steps safely (min 1)
+DEFAULT_SAVE_STEPS=$((8192 / GPU_COUNT))
+(( DEFAULT_SAVE_STEPS < 1 )) && DEFAULT_SAVE_STEPS=1

modelopt/torch/export/plugins/hf_spec_export.py (6)

18-20: Record transformers version (best-effort) for traceability.

Populate transformers_version to aid downstream tooling.

@@
-import torch
+import torch
 import torch.nn as nn
+try:
+    import transformers as _hf
+    _TRANSFORMERS_VERSION = getattr(_hf, "__version__", None)
+except Exception:
+    _TRANSFORMERS_VERSION = None

---

76-79: Harden lm_head fallback; raise if neither drafter nor base provides it.

Prevents silent KeyError from model.state_dict()["lm_head.weight"].

@@
-    if "eagle_lm_head.weight" not in eagle_state:
-        export_state_dict["lm_head.weight"] = model.state_dict()["lm_head.weight"]
+    if "eagle_lm_head.weight" not in eagle_state:
+        base_state = model.state_dict()
+        if "lm_head.weight" in base_state:
+            export_state_dict["lm_head.weight"] = base_state["lm_head.weight"]
+        else:
+            raise KeyError("Missing 'eagle_lm_head.weight' in drafter and 'lm_head.weight' in base model.")

---

127-134: Guard eagle_config/config access before getattr.

Avoids AttributeError when either container is missing.

@@
-    def _get_config_from_eagle_config_or_base_config(key: str, model: nn.Module):
-        if getattr(model.eagle_config, key, None) is not None:
-            return getattr(model.eagle_config, key)
-        elif getattr(model.config, key, None) is not None:
-            return getattr(model.config, key)
-        else:
-            return None
+    def _get_config_from_eagle_config_or_base_config(key: str, model: nn.Module):
+        eagle_cfg = getattr(model, "eagle_config", None)
+        if eagle_cfg is not None and getattr(eagle_cfg, key, None) is not None:
+            return getattr(eagle_cfg, key)
+        base_cfg = getattr(model, "config", None)
+        if base_cfg is not None and getattr(base_cfg, key, None) is not None:
+            return getattr(base_cfg, key)
+        return None

---

89-108: Preserve original config fields and merge; set transformers_version.

Avoid dropping unknown config keys and persist the detected Transformers version.

@@
-        "transformers_version": None,
+        "transformers_version": _TRANSFORMERS_VERSION,
@@
-    return template_config
+    # Merge with original to preserve unknown keys.
+    merged = {**config_data, **template_config}
+    merged["eagle_config"] = {
+        **config_data.get("eagle_config", {}),
+        **template_config["eagle_config"],
+    }
+    return merged

Also applies to: 149-149

---

84-88: set_config_if_spec_decoding crashes on vanilla HF models.

Direct len(model._modelopt_state) will raise when the attribute is absent.

@@
-    if len(model._modelopt_state) != 1 or model._modelopt_state[0][0] != "eagle":
+    opt_modes = getattr(model, "_modelopt_state", None)
+    if (
+        not isinstance(opt_modes, (list, tuple))
+        or len(opt_modes) != 1
+        or opt_modes[0][0] != "eagle"
+    ):
         # return as is
         return config_data

---

51-65: Guard for missing eagle_module to avoid AttributeError.

Accessing model.eagle_module without a presence check can crash on non-Eagle models.

@@
-    # Check if the state dict keys match
-    _check_state_dict_keys_match(model.eagle_module, EAGLE_MODELOPT_TO_OFFICIAL["required"])
+    # Ensure eagle_module exists
+    if not hasattr(model, "eagle_module"):
+        return post_state_dict
+    # Check if the state dict keys match
+    _check_state_dict_keys_match(model.eagle_module, EAGLE_MODELOPT_TO_OFFICIAL["required"])

modelopt/torch/export/unified_export_hf.py (4)

494-498: Fix _quant_applied logic; make it reliable.

Current check can return True for empty configs and is used nowhere. Tighten and use it.

@@
-def _quant_applied(hf_quant_config: dict) -> bool:
-    """Check if any quantization is applied."""
-    q = hf_quant_config.get("quantization", {})
-    return not (q.get("quant_algo") == QUANTIZATION_NONE and not q.get("quantized_layers"))
+def _quant_applied(hf_quant_config: dict) -> bool:
+    """Return True iff any quantization is configured."""
+    q = hf_quant_config.get("quantization") or {}
+    algo = q.get("quant_algo")
+    layers = q.get("quantized_layers")
+    return (algo is not None and algo != QUANTIZATION_NONE) or bool(layers)

---

519-523: Don’t emit empty hf_quant_config.json.

Only write when quantization is actually applied; also use Path for portability.

@@
-        # Save hf_quant_config.json for backward compatibility
-        with open(f"{export_dir}/hf_quant_config.json", "w") as file:
-            json.dump(hf_quant_config, file, indent=4)
+        # Save hf_quant_config.json only when quantization is applied
+        if _quant_applied(hf_quant_config):
+            with open(Path(export_dir) / "hf_quant_config.json", "w") as file:
+                json.dump(hf_quant_config, file, indent=4)

---

526-527: Guard spec-decoding post-state transform when ModelOpt state is absent.

Prevents AttributeError on vanilla HF models.

@@
-        post_state_dict = rename_and_prune_if_spec_decoding(model, post_state_dict)
+        if hasattr(model, "_modelopt_state"):
+            post_state_dict = rename_and_prune_if_spec_decoding(model, post_state_dict)

---

539-544: Preserve quantization_config and avoid spec-config clobbering; gate on quant applied.

Move assignment after spec-config transform and gate it. Also remove stale key when no-quant.

@@
-        config_data["quantization_config"] = hf_quant_config
-
-        config_data = set_config_if_spec_decoding(model, config_data)
+        if hasattr(model, "_modelopt_state"):
+            config_data = set_config_if_spec_decoding(model, config_data)
+        if _quant_applied(hf_quant_config):
+            config_data["quantization_config"] = hf_quant_config
+        else:
+            config_data.pop("quantization_config", None)

examples/speculative_decoding/launch_train.sh (4)

111-116: Quote config path to handle spaces and special chars.

Unquoted paths can break argument parsing.

-    SPECULATIVE_ARGS="--eagle_config $EAGLE_CONFIG"
+    SPECULATIVE_ARGS="--eagle_config \"$EAGLE_CONFIG\""

---

130-153: Quote key CLI values to be robust to spaces.

Quote MODEL, OUTPUT_DIR, and DATA in the assembled command.

-    --model_name_or_path $MODEL \
+    --model_name_or_path \"$MODEL\" \
@@
-    --output_dir $OUTPUT_DIR \
+    --output_dir \"$OUTPUT_DIR\" \
@@
-    --data_path $DATA \
+    --data_path \"$DATA\" \

---

103-106: Remove or wire unused variables.

REDRAFTER_TOKENS and REDRAFTER_NUM_LAYERS are set but unused.

---

77-80: Improve invalid-argument error message.

Print the full offending token; current expansion may mangle it.

-      >&2 printf "Error: Invalid argument ${1#*=}\n"
+      >&2 printf "Error: Invalid argument: %s\n" "$1"

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.359Z
Learning: In the TensorRT-Model-Optimizer codebase, EAGLE1_DEFAULT_CFG in config.py already uses deepcopy(default_eagle_config) in its definition, making additional deepcopy calls unnecessary when using ALGO_TO_CONFIG mapping in tests.

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.359Z
Learning: In the TensorRT-Model-Optimizer codebase, both EAGLE1_DEFAULT_CFG and EAGLE3_DEFAULT_CFG in config.py already use deepcopied configurations (either directly via deepcopy() or through variables that were already deepcopied), making additional deepcopy calls unnecessary when using ALGO_TO_CONFIG mapping in tests.

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.359Z
Learning: In the TensorRT-Model-Optimizer codebase, the EAGLE1_DEFAULT_CFG and EAGLE3_DEFAULT_CFG configurations in config.py are already deepcopied, so additional deepcopy calls are not needed when using ALGO_TO_CONFIG mapping in tests.

examples/speculative_decoding/server_generate.py (1)

49-49: Fix boolean CLI parsing for --chat.

type=bool treats any non-empty string (even "False") as True. Use BooleanOptionalAction so --chat/--no-chat work correctly.

-parser.add_argument("--chat", default=True, type=bool, help="Use chat mode")
+parser.add_argument(
+    "--chat",
+    action=argparse.BooleanOptionalAction,
+    default=True,
+    help="Use chat mode (use --no-chat to disable)",
+)

modelopt/torch/export/plugins/hf_spec_export.py (4)

89-125: Preserve original config_data fields and set transformers_version; deep-merge eagle_config.

Current code replaces the entire config, dropping unknown fields. Merge defaults into the original and record the Transformers version.

@@
-    # This is the config keys in official checkpoint.
+    # This is the config keys in official checkpoint.
     template_config = {
@@
-        "transformers_version": None,
+        "transformers_version": getattr(transformers, "__version__", None),
@@
-    for key in template_config:
+    for key in template_config:
         value = template_config[key]
@@
-            template_config[key] = new_value
+            template_config[key] = new_value
@@
-    return template_config
+    # Merge: keep any unknown keys from original config_data
+    merged = {**config_data, **template_config}
+    # Deep-merge nested eagle_config
+    merged["eagle_config"] = {
+        **config_data.get("eagle_config", {}),
+        **template_config["eagle_config"],
+    }
+    return merged

Also applies to: 135-149

---

76-79: Harden fallback for lm_head.weight export.

Explicitly check base model state and raise a clear error if neither key exists.

-    # TODO: (hg) this is a temp fix. Find cleaner way to do this.
-    if "eagle_lm_head.weight" not in eagle_state:
-        export_state_dict["lm_head.weight"] = model.state_dict()["lm_head.weight"]
+    # TODO: (hg) this is a temp fix. Find cleaner way to do this.
+    if "eagle_lm_head.weight" not in eagle_state:
+        base_state = model.state_dict()
+        if "lm_head.weight" in base_state:
+            export_state_dict["lm_head.weight"] = base_state["lm_head.weight"]
+        else:
+            raise KeyError(
+                "Missing 'eagle_lm_head.weight' in draft and 'lm_head.weight' in base model."
+            )

---

127-134: Guard access to eagle_config/config to avoid AttributeError.

-    def _get_config_from_eagle_config_or_base_config(key: str, model: nn.Module):
-        if getattr(model.eagle_config, key, None) is not None:
-            return getattr(model.eagle_config, key)
-        elif getattr(model.config, key, None) is not None:
-            return getattr(model.config, key)
-        else:
-            return None
+    def _get_config_from_eagle_config_or_base_config(key: str, model: nn.Module):
+        eagle_cfg = getattr(model, "eagle_config", None)
+        if eagle_cfg is not None and getattr(eagle_cfg, key, None) is not None:
+            return getattr(eagle_cfg, key)
+        base_cfg = getattr(model, "config", None)
+        if base_cfg is not None and getattr(base_cfg, key, None) is not None:
+            return getattr(base_cfg, key)
+        return None

---

51-66: Guard for Eagle mode and presence of eagle_module before accessing it.

If _modelopt_state indicates Eagle but eagle_module is missing, this raises AttributeError. Early-return safely.

 def rename_and_prune_if_spec_decoding(model: nn.Module, post_state_dict: dict):
@@
-    opt_modes = getattr(model, "_modelopt_state", None)
+    opt_modes = getattr(model, "_modelopt_state", None)
     if (
         not isinstance(opt_modes, (list, tuple))
         or len(opt_modes) != 1
         or opt_modes[0][0] != "eagle"
     ):
         # if there's other opts, return as is
         return post_state_dict
+    if not hasattr(model, "eagle_module"):
+        # Not an Eagle-wrapped model; return unchanged
+        return post_state_dict
@@
-    _check_state_dict_keys_match(model.eagle_module, EAGLE_MODELOPT_TO_OFFICIAL["required"])
+    _check_state_dict_keys_match(model.eagle_module, EAGLE_MODELOPT_TO_OFFICIAL["required"])

examples/speculative_decoding/README.md (1)

28-33: Provide concrete install commands (replace placeholder).

Make this copy-pastable; include PyPI install and editable source install with HF extras.

-Install Modelopt with `hf` dependencies and other requirements for this example:
+Install ModelOpt with `hf` extras and example requirements:
@@
-```bash
-pip install -e ...
-pip install -r requirements.txt
-```
+```bash
+# Option A: Install published package (recommended)
+pip install "nvidia-modelopt[hf]"
+
+# Option B: Install from source
+pip install -e .[hf]
+
+# Example-specific deps
+pip install -r requirements.txt
+```

examples/speculative_decoding/server_generate.py (3)

56-57: Use list default for --system_prompt to avoid odd joins.

Defaulting to "" makes " ".join(args.system_prompt) iterate characters. Use an empty list.

-parser.add_argument("--system_prompt", nargs="+", type=str, default="", help="System prompt")
+parser.add_argument("--system_prompt", nargs="+", type=str, default=[], help="System prompt")

---

43-44: Lower default thread count to a safer value.

256 workers can easily overwhelm local servers and hit rate limits. Suggest 32–64 by default.

-parser.add_argument(
-    "--num_threads", type=int, default=256, help="Number of threads to use (batch size)"
+parser.add_argument(
+    "--num_threads", type=int, default=64, help="Number of threads to use (batch size)"
 )

---

60-65: Use context manager for JSON load for consistency and file handle safety.

-if args.data_path.endswith("jsonl"):
-    with open(args.data_path) as f:
-        data = [json.loads(line) for line in f]
-else:
-    data = json.load(open(args.data_path))
+if args.data_path.endswith("jsonl"):
+    with open(args.data_path) as f:
+        data = [json.loads(line) for line in f]
+else:
+    with open(args.data_path) as f:
+        data = json.load(f)

modelopt/torch/export/plugins/hf_spec_export.py (1)

18-20: Record Transformers version for provenance.

Import transformers to populate transformers_version in the exported config.

 import torch
 import torch.nn as nn
+import transformers

examples/speculative_decoding/README.md (2)

47-48: Brand consistency: “ModelOpt” capitalization.

Use “ModelOpt” consistently.

-This one-line command runs a minimal example workflow of training and exporting an EAGLE draft model in Modelopt. Specifically, it
+This one-line command runs a minimal example workflow of training and exporting an EAGLE draft model in ModelOpt. Specifically, it

---

229-236: Validate or trim the Support Matrix to avoid drift.

Section can get stale quickly; either confirm the list against current HF path support or drop the table and link to dynamic docs.

Would you like me to open a follow-up PR to auto-generate this table from a tested registry?

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro